# 导入工具包
import torch

# 设定模型源source
source = 'huggingface/pytorch-transformers'
# 设定模型的名称
model_name = 'ItcastAI/bert_cn_finetunning'

# 加载映射器和模型
tokenizer = torch.hub.load(source, 'tokenizer', model_name)
model = torch.hub.load(source, 'modelForSequenceClassification', model_name)


# 编写得到标签的函数
def get_label(text):
    # 第一步对中文文本进行预训练模型的编码
    index = tokenizer.encode(text)
    # 将编码结果封装成torch.tensor类型
    tokens_tensor = torch.tensor([index])
    # 进入评估模式
    with torch.no_grad():
        result = model(tokens_tensor)
    # 将最大值所代表的标签作为模型的预测标签
    predicted_label = torch.argmax(result[0]).item()
    return predicted_label


if __name__ == '__main__':
    # 初始化一段中文文本
    text = "早餐不好,服务不到位,晚餐无西餐,早餐晚餐相同,房间条件不好"
    print("输入文本为:", text)
    print("预测标签为:", get_label(text))
